/*
 * Copyright (c) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef SMART_SERVE_PLUGIN_TASK_H
#define SMART_SERVE_PLUGIN_TASK_H

#include <atomic>
#include <condition_variable>
#include <mutex>
#include <string>

namespace OHOS {
namespace SmartServe {

class EvalTask {
public:
    enum class Type { LLM };
    enum class TaskState { UNDONE, FINISHED, ABORTED, CANCELLED };
    TaskState state_{EvalTask::TaskState::UNDONE};

    bool isDone() const // 标识session不再会访问task
    {
        return state_ == TaskState::FINISHED || state_ == TaskState::ABORTED || state_ == TaskState::CANCELLED;
    }
    inline std::string doneReason()
    {
        switch (state_) {
            case TaskState::CANCELLED:
                return "Cancelled";
            case TaskState::FINISHED:
                return "Finished";
            case TaskState::ABORTED:
                return "Aborted";
            default:
                return "";
        }
    }

    EvalTask(Type type) : type_(type)
    {}
    virtual ~EvalTask() = default;

    void Done(TaskState state)
    {
        {
            std::unique_lock<std::mutex> lock(mutex_);
            state_ = state;
        }
        cv_.notify_all();
    }

    Type type_;

    std::mutex mutex_;
    std::condition_variable cv_;
};

class EvalLLMTask : public EvalTask {
public:
    EvalLLMTask(pid_t pid, const std::string& modelName, const std::string& inputs)
        : EvalTask(EvalTask::Type::LLM), pid_(pid), modelName_(modelName), inputs_(inputs)
    {}

    void AppendResult(const std::string& chunk) // 内部加 Task 锁
    {
        {
            std::unique_lock<std::mutex> lock(mutex_);
            results_ += chunk;
        }
        cv_.notify_all();
    }

    void WaitForRemainingResult(std::unique_lock<std::mutex>& lock) // 要求加 Task 锁
    {
        cv_.wait(lock, [this]() { return nGot_ < results_.size() || isDone(); });
    }

    std::string GetRemainingResult() // 要求加 Task 锁
    {
        auto ret = results_.substr(nGot_);
        nGot_ = results_.size();
        return ret;
    }

    pid_t pid_;
    std::string modelName_;
    std::string inputs_;

    std::string results_;
    std::size_t nGot_ = 0;
};

} // namespace SmartServe
} // namespace OHOS

#endif // SMART_SERVE_PLUGIN_TASK_H
